import torch  # PyTouch 深度学习框架

inputs = torch.tensor(
    [[0.43, 0.15, 0.89],  # Your     (x^1)
     [0.55, 0.87, 0.66],  # journey  (x^2)
     [0.57, 0.85, 0.64],  # starts   (x^3)
     [0.22, 0.58, 0.33],  # with     (x^4)
     [0.77, 0.25, 0.10],  # one      (x^5)
     [0.05, 0.80, 0.55]]  # step     (x^6)
)

# 1.计算注意力分数
query = inputs[1]
print(query)
# tensor([0.5500, 0.8700, 0.6600])
attn_scores_2 = torch.empty(inputs.shape[0])
# 以第2行词元journey的词向量为例，与每一行的向量做点积，计算出中间值w，即注意力分数
for i, x_i in enumerate(inputs):
    attn_scores_2[i] = torch.dot(x_i, query)
print(attn_scores_2)
# tensor([0.9544, 1.4950, 1.4754, 0.8434, 0.7070, 1.0865])


# 2.对注意力分数进行归一化处理
attn_weights_2_tmp = attn_scores_2 / attn_scores_2.sum()
print("Attention weights:", attn_weights_2_tmp)
print("Sum:", attn_weights_2_tmp.sum())
# Attention weights: tensor([0.1455, 0.2278, 0.2249, 0.1285, 0.1077, 0.1656])
# Sum: tensor(1.0000)


# 在实际应用中，使用softmax函数进行归一化更为常见，而且是一种更可取的做法。
# 这种方法更好地处理了极值，并在训练期间提供了更有利的梯度特性。
# 以下是用于归一化注意力分数softmax函数的基础实现：
def softmax_naive(x):
    return torch.exp(x) / torch.exp(x).sum(dim=0)


attn_weights_2_naive = softmax_naive(attn_scores_2)
print("Attention weights:", attn_weights_2_naive)
print("Sum:", attn_weights_2_naive.sum())
# Attention weights: tensor([0.1385, 0.2379, 0.2333, 0.1240, 0.1082, 0.1581])
# Sum: tensor(1.)
# 另外，softmax函数可以保证注意力权重总是正值，这使得输出可以被解释为概率或相对重要性，其中权重越高表示重要程度越高。


# 3.计算最终的上下文向量
query = inputs[1]
context_vec_2 = torch.zeros(inputs.shape)
print(inputs.shape)
# 以第2行词元journey的词向量为例，与每一行的向量做点积，计算出中间值w，即注意力分数
for i, x_i in enumerate(inputs):
    context_vec_2 += attn_weights_2_naive[i] * x_i
print(context_vec_2)
# tensor([0.4419, 0.6515, 0.5683])


# ---------------------------------------------------------------------------

# 点积的计算公式示例：
res = 0.
for idx, element in enumerate(inputs[0]):
    res += inputs[0][idx] * query[idx]
print(res)
print(torch.dot(inputs[0], query))
# tensor(0.9544)
# tensor(0.9544)
